package priority_channel

import (
	"container/heap"
	"context"
	"sync"
	"time"
)

// Priority levels for events
const (
	PriorityHigh   = 0
	PriorityMedium = 1
	PriorityLow    = 2
)

// Event represents a message in the priority channel
type Event struct {
	ID         string      `json:"id"`
	Priority   int         `json:"priority"`
	Type       string      `json:"type"`
	Data       interface{} `json:"data"`
	Timestamp  time.Time   `json:"timestamp"`
	RetryCount int         `json:"retry_count"`
	MaxRetries int         `json:"max_retries"`
}

// PriorityChannel interface defines the contract for priority channel operations
type PriorityChannel interface {
	// Push adds an event to the channel
	Push(event *Event) error

	// Pop removes and returns the highest priority event (non-blocking)
	Pop() (*Event, error)

	// Receive blocks and waits for the next event with select context
	Receive(ctx context.Context) (*Event, error)

	// Peek returns the highest priority event without removing it
	Peek() (*Event, error)

	// Size returns the current number of events in the channel
	Size() int

	// IsEmpty returns true if the channel is empty
	IsEmpty() bool

	// GetMemoryUsage returns the current memory usage in bytes
	GetMemoryUsage() int64

	// Close gracefully shuts down the channel
	Close() error
}

// MemoryMonitor interface for monitoring memory usage
type MemoryMonitor interface {
	GetMemoryUsage() int64
	GetMaxMemoryUsage() int64
	ResetMaxMemoryUsage()
	updateMemoryUsage(eventCount int)
}

// priorityChannel implements the PriorityChannel interface
type priorityChannel struct {
	mu            sync.RWMutex
	heap          *eventHeap
	receiveChan   chan *Event
	ctx           context.Context
	cancel        context.CancelFunc
	memoryMonitor MemoryMonitor
	config        *ChannelConfig
}

// ChannelConfig holds configuration for the priority channel
type ChannelConfig struct {
	BufferSize     int           `json:"buffer_size"`
	MaxRetries     int           `json:"max_retries"`
	RetryDelay     time.Duration `json:"retry_delay"`
	MemoryLimit    int64         `json:"memory_limit"`
	PriorityLevels []int         `json:"priority_levels"`
}

// DefaultChannelConfig returns default configuration
func DefaultChannelConfig() *ChannelConfig {
	return &ChannelConfig{
		BufferSize:     1000,
		MaxRetries:     3,
		RetryDelay:     time.Second,
		MemoryLimit:    1024 * 1024 * 100, // 100MB
		PriorityLevels: []int{PriorityHigh, PriorityMedium, PriorityLow},
	}
}

// eventHeap implements heap.Interface for the priority channel
type eventHeap []*Event

func (h eventHeap) Len() int { return len(h) }

func (h eventHeap) Less(i, j int) bool {
	// Lower priority number means higher priority
	if h[i].Priority != h[j].Priority {
		return h[i].Priority < h[j].Priority
	}
	// If priorities are equal, older events have higher priority
	return h[i].Timestamp.Before(h[j].Timestamp)
}

func (h eventHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *eventHeap) Push(x interface{}) {
	*h = append(*h, x.(*Event))
}

func (h *eventHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// NewPriorityChannel creates a new priority channel with the specified configuration
func NewPriorityChannel(config *ChannelConfig) PriorityChannel {
	if config == nil {
		config = DefaultChannelConfig()
	}

	ctx, cancel := context.WithCancel(context.Background())

	pc := &priorityChannel{
		heap:          &eventHeap{},
		receiveChan:   make(chan *Event, config.BufferSize),
		ctx:           ctx,
		cancel:        cancel,
		memoryMonitor: newMemoryMonitor(),
		config:        config,
	}

	// Initialize the heap
	heap.Init(pc.heap)

	return pc
}

// Push adds an event to the channel
func (pc *priorityChannel) Push(event *Event) error {
	pc.mu.Lock()
	defer pc.mu.Unlock()

	select {
	case <-pc.ctx.Done():
		return ErrChannelClosed
	default:
		// Check memory limit
		if pc.memoryMonitor.GetMemoryUsage() > pc.config.MemoryLimit {
			return ErrMemoryLimitExceeded
		}

		heap.Push(pc.heap, event)
		pc.memoryMonitor.updateMemoryUsage(pc.heap.Len())

		// Try to send to receive channel (non-blocking)
		select {
		case pc.receiveChan <- event:
		default:
			// Channel is full, event is still in heap
		}

		return nil
	}
}

// Pop removes and returns the highest priority event (non-blocking)
func (pc *priorityChannel) Pop() (*Event, error) {
	pc.mu.Lock()
	defer pc.mu.Unlock()

	select {
	case <-pc.ctx.Done():
		return nil, ErrChannelClosed
	default:
		if pc.heap.Len() > 0 {
			event := heap.Pop(pc.heap).(*Event)
			pc.memoryMonitor.updateMemoryUsage(pc.heap.Len())
			return event, nil
		}
		return nil, ErrChannelEmpty
	}
}

// Receive blocks and waits for the next event with select context
func (pc *priorityChannel) Receive(ctx context.Context) (*Event, error) {
	// First try to get from heap (non-blocking)
	pc.mu.Lock()
	if pc.heap.Len() > 0 {
		event := heap.Pop(pc.heap).(*Event)
		pc.memoryMonitor.updateMemoryUsage(pc.heap.Len())
		pc.mu.Unlock()
		return event, nil
	}
	pc.mu.Unlock()

	// If heap is empty, wait for new events
	select {
	case <-pc.ctx.Done():
		return nil, ErrChannelClosed
	case <-ctx.Done():
		return nil, ctx.Err()
	case event := <-pc.receiveChan:
		// Remove from heap since it was sent to receive channel
		pc.mu.Lock()
		// Find and remove the event from heap
		for i, e := range *pc.heap {
			if e.ID == event.ID {
				heap.Remove(pc.heap, i)
				break
			}
		}
		pc.memoryMonitor.updateMemoryUsage(pc.heap.Len())
		pc.mu.Unlock()
		return event, nil
	}
}

// Peek returns the highest priority event without removing it
func (pc *priorityChannel) Peek() (*Event, error) {
	pc.mu.RLock()
	defer pc.mu.RUnlock()

	select {
	case <-pc.ctx.Done():
		return nil, ErrChannelClosed
	default:
		if pc.heap.Len() > 0 {
			return (*pc.heap)[0], nil
		}
		return nil, ErrChannelEmpty
	}
}

// Size returns the current number of events in the channel
func (pc *priorityChannel) Size() int {
	pc.mu.RLock()
	defer pc.mu.RUnlock()

	select {
	case <-pc.ctx.Done():
		return 0
	default:
		return pc.heap.Len()
	}
}

// IsEmpty returns true if the channel is empty
func (pc *priorityChannel) IsEmpty() bool {
	pc.mu.RLock()
	defer pc.mu.RUnlock()

	select {
	case <-pc.ctx.Done():
		return true
	default:
		return pc.heap.Len() == 0
	}
}

// GetMemoryUsage returns the current memory usage in bytes
func (pc *priorityChannel) GetMemoryUsage() int64 {
	pc.mu.RLock()
	defer pc.mu.RUnlock()

	select {
	case <-pc.ctx.Done():
		return 0
	default:
		return pc.memoryMonitor.GetMemoryUsage()
	}
}

// Close gracefully shuts down the channel
func (pc *priorityChannel) Close() error {
	pc.mu.Lock()
	defer pc.mu.Unlock()

	select {
	case <-pc.ctx.Done():
		return ErrChannelClosed
	default:
		pc.cancel()
		close(pc.receiveChan)
		return nil
	}
}
